import os
import re
import sys

class DataStruct:
    def __init__(self,attr_name,start_addr,size,lib_name,file_name):
        self.attr = attr_name
        self.start = start_addr
        self.size = size
        self.lib = lib_name
        self.file = file_name
            
class ParseMap:                
    def __init__(self,file,strategy = 1):
        self.NL = "\r\n"
        self.map = {}
        self.line_num = 0
        self.fd = open(file)
        self.removeHeader(self.fd)

        self.strategy = strategy
        self.processTextSection(self.fd)        

    def removeHeader(self,f):        
        while True:
            line = f.readline()
            self.line_num +=1
            if line == "":
                break
                
            if line.startswith("Linker script and memory map"):
                f.readline()                
                self.line_num +=1
                while not f.readline().startswith(" "):
                    self.line_num +=1
                break

    def processOne(self,map,fd,sline):
        sline[0] = sline[0][0:20]
        attr_name=""
        start_addr=""            
        libname = ""
        filename = ""
        size = 0

        comp = re.compile(r'([^(]*)\(([^\)]*)')
        m = comp.match(sline[3])
        if m:
            libname = m.group(1)
            filename= m.group(2)
        else:
            libname=sline[3]
            filename=""

        attr_name = sline[0]
        start_addr = sline[1]
        size = int(sline[2],16)
        
        key = "" 
        if self.strategy == 1:
            key = sline[0][0:4] + libname
            if key in map:
                map[key].size += size
            else:
                map[key] = DataStruct(attr_name,start_addr,size,libname,filename)
                
    def processTextSection(self,fd):
        NL= self.NL                        
        map = {}
        
        while True:
            line=fd.readline()
            self.line_num +=1
            
            if line == "":                
                break

            sline = line.strip().split()  
            if len(sline) > 0 and sline[0].startswith(".debug_"):
                break

            doit = False
            if len(sline) == 4 and sline[2].startswith("0x"):           
                doit = True    
            elif len(sline) == 1 and sline[0]:
                s0 = sline[0]
                if s0.startswith(".tex") or s0.startswith(".bss") or s0.startswith(".roda"):
                    oldline = line
                    line = fd.readline()
                    self.line_num +=1
                    sline = sline + line.split()
                    doit = True
                    
            if doit:
                self.processOne(map,fd,sline)
                
        self.map = map        

    def print(self):
        text_size = 0
        ro_size = 0
        bss_size = 0
        category=1

        tsize = 0        
        for data in self.map.values():
            if category == 1 and data.attr.startswith(".rodata"):
                print("")
                category =2
            elif category == 2 and data.attr.startswith(".bss"):
                print("")
                category = 3
                
            if category == 1:
                text_size += data.size
                tsize = text_size
            elif category == 2:
                ro_size += data.size
                tsize = ro_size
            else:
                bss_size += data.size
                tsize = bss_size
                
            print("%-20s"%data.attr, "%-20s"%data.start,"0x%-10x"%data.size,"%dKB"%(tsize/1024),data.lib[-50:])

        print("")
        print("text_size %dKB, rodata %dKB, bss %dKB, total %dKB"%(text_size/1024,ro_size/1024,bss_size/1024,(text_size+ro_size+bss_size)/1024))


    def printDataStruct(self,data):  
        print("%-20s"%data.attr, "%-20s"%data.start, "%5dKB"%(data.size/1024),"%6d"%data.size,data.lib[-50:])
        
    def showTop(self,att,number):
        lst = []
        for data in self.map.values():
            if data.attr.startswith(att):
                lst.append(data)
                
        lst.sort(key=lambda data:data.size ,reverse = True)

        for  i in range(0, min(len(lst),number)):
            self.printDataStruct(lst[i])
            
    def close(self):
        self.fd.close()
                
if __name__ == "__main__":                
    map = ParseMap(r"E:\dataset\watch\software\bootloader\Debug\watch_bootloader.map")
    print("")
    if len(sys.argv) == 3:    
        map.showTop(sys.argv[1],int(sys.argv[2]))
    else:
        map.print()
        
    map.close()